# 更新时间：2025年9月24日 20:33:49 亲测可用

import os
import re
import tkinter as tk
from tkinter import simpledialog, filedialog, messagebox

def get_user_input():
    """获取用户输入的目录和标签"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    # 选择目录
    directory = filedialog.askdirectory(title="选择包含Markdown文件的目录")
    if not directory:
        return None, None
    
    # 输入标签
    tags = simpledialog.askstring("输入标签", "请输入要添加的标签（多个标签用逗号分隔）:")
    if not tags:
        return directory, None
    
    # 清理标签（去除空格，分割成列表）
    tag_list = [tag.strip() for tag in tags.split(',') if tag.strip()]
    
    return directory, tag_list

def extract_front_matter(content):
    """从Markdown内容中提取front matter（如果有的话）"""
    front_matter = {}
    if content.startswith('---'):
        end_index = content.find('---', 3)
        if end_index != -1:
            front_matter_text = content[3:end_index].strip()
            # 解析front matter
            for line in front_matter_text.split('\n'):
                if ':' in line:
                    key, value = line.split(':', 1)
                    key = key.strip()
                    value = value.strip()
                    
                    # 处理tags字段
                    if key.lower() == 'tags':
                        if value.startswith('[') and value.endswith(']'):
                            # 列表格式的tags
                            tags_str = value[1:-1]
                            tags = [tag.strip().strip('"\'') for tag in tags_str.split(',') if tag.strip()]
                            front_matter['tags'] = tags
                        else:
                            # 单个标签或空格分隔的标签
                            tags = [tag.strip().strip('"\'') for tag in value.split() if tag.strip()]
                            front_matter['tags'] = tags
                    else:
                        front_matter[key] = value
            
            return front_matter, content[end_index+3:].lstrip()
    
    return front_matter, content

def add_tags_to_file(file_path, tags_to_add):
    """为单个Markdown文件添加标签"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件 {file_path} 时出错: {e}")
        return False
    
    front_matter, remaining_content = extract_front_matter(content)
    
    # 获取现有标签
    existing_tags = front_matter.get('tags', [])
    if not isinstance(existing_tags, list):
        existing_tags = [existing_tags]
    
    # 过滤掉已存在的标签
    new_tags = [tag for tag in tags_to_add if tag not in existing_tags]
    
    if not new_tags:
        print(f"文件 {file_path} 已包含所有指定标签，跳过")
        return True
    
    # 添加新标签
    updated_tags = existing_tags + new_tags
    
    # 更新front matter
    front_matter['tags'] = updated_tags
    
    # 重新构建front matter
    front_matter_lines = ['---']
    for key, value in front_matter.items():
        if key == 'tags':
            # 将tags格式化为YAML列表
            if value:
                front_matter_lines.append(f'tags: [{", ".join([f"\"{tag}\"" for tag in value])}]')
            else:
                front_matter_lines.append('tags: []')
        else:
            front_matter_lines.append(f'{key}: {value}')
    front_matter_lines.append('---')
    
    # 重新构建文件内容
    new_content = '\n'.join(front_matter_lines) + '\n\n' + remaining_content
    
    # 写回文件
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(new_content)
        print(f"成功为 {file_path} 添加标签: {', '.join(new_tags)}")
        return True
    except Exception as e:
        print(f"写入文件 {file_path} 时出错: {e}")
        return False

def process_markdown_files(directory, tags_to_add):
    """处理目录中的所有Markdown文件"""
    if not tags_to_add:
        print("没有提供标签，退出")
        return
    
    md_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.md'):
                md_files.append(os.path.join(root, file))
    
    if not md_files:
        print(f"在目录 {directory} 中未找到Markdown文件")
        return
    
    print(f"找到 {len(md_files)} 个Markdown文件")
    
    success_count = 0
    for file_path in md_files:
        if add_tags_to_file(file_path, tags_to_add):
            success_count += 1
    
    print(f"处理完成。成功处理 {success_count}/{len(md_files)} 个文件")

def main():
    """主函数"""
    directory, tags_to_add = get_user_input()
    
    if not directory:
        print("未选择目录，退出")
        return
    
    if not tags_to_add:
        print("未提供标签，退出")
        return
    
    process_markdown_files(directory, tags_to_add)
    
    # 显示完成消息
    root = tk.Tk()
    root.withdraw()
    messagebox.showinfo("完成", f"已为目录 {directory} 中的所有Markdown文件添加标签: {', '.join(tags_to_add)}")

if __name__ == "__main__":
    main() 